{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 导包，加载图片，组成4D的张量\n",
    "\n",
    "- 图片放在`data`目录下，使用`img_to_array()`加载图片为numpy arrays\n",
    "- 将图片按batch组合，同时处理多张图片，扩充程序的扩展性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 400, 400, 3)\n",
      "Input Img Shape: (2, 400, 400, 3)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "from utils import  img_to_array,array_to_img,visualize_grid,view_images\n",
    "\n",
    "# params\n",
    "DIMS = (400, 400)\n",
    "CAT1 = 'cat1.jpg'\n",
    "CAT2 = 'cat2.jpg'\n",
    "data_path = './data/'\n",
    "\n",
    "# 加载两张小猫图片\n",
    "img1 = img_to_array(data_path + CAT1, DIMS)\n",
    "img2 = img_to_array(data_path + CAT2, DIMS, view=True)\n",
    "\n",
    "# 联合两张图片到一个batch,shape为(2, 400, 400, 3)\n",
    "input_img = np.concatenate([img1, img2], axis=0)\n",
    "\n",
    "print(img2.shape)\n",
    "print(\"Input Img Shape: {}\".format(input_img.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 构建仿射变换生成矩阵和双线性插值函数\n",
    "\n",
    "`affine_grid_generator(height, width, M)`为仿射变换函数:\n",
    "\n",
    "- height,width,为图片长和宽\n",
    "- M： 仿射矩阵。 shape为(num_batch, 2, 3). \n",
    "- Return: 返回sampling grid。shape为(num_batch, H, W, 2)\n",
    "\n",
    "`bilinear_sampler(input_img, x, y)`为双线性采样：\n",
    "\n",
    "- input_img：采样的原始图片，(B, H, W, C)\n",
    "- x,y: affine_grid_generator输出\n",
    "- Return：采样后的值"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def affine_grid_generator(height, width, M):\n",
    "    num_batch = M.shape[0] # 获取batchsize\n",
    "\n",
    "    # 创建棋盘grid,平分整个图片\n",
    "    x = np.linspace(-1, 1, width)\n",
    "    y = np.linspace(-1, 1, height)\n",
    "    x_t, y_t = np.meshgrid(x, y)\n",
    "\n",
    "    # 制作列向量(xt, yt, 1) \n",
    "    ones = np.ones(np.prod(x_t.shape))\n",
    "    sampling_grid = np.vstack([x_t.flatten(), y_t.flatten(), ones])\n",
    "    # 把所有像素整到一起  sampling_grid的shape为(batch, 3, H*W)  \n",
    "    sampling_grid = np.resize(sampling_grid, (num_batch, 3, height*width)) # 列向量扩展batch\n",
    "    \n",
    "    # 做仿射矩阵运算 M*K\n",
    "    batch_grids = np.matmul(M, sampling_grid)\n",
    "    # batch grid 的 shape (num_batch, 2, H*W)\n",
    "    \n",
    "    # 把H,W拆开\n",
    "    batch_grids = batch_grids.reshape(num_batch, 2, height, width)\n",
    "    batch_grids = np.moveaxis(batch_grids, 1, -1) #调整为(num_batch, H, W, 2)\n",
    "\n",
    "    # sanity check\n",
    "    print(\"Transformation Matrices: {}\".format(M.shape))\n",
    "    print(\"Sampling Grid: {}\".format(sampling_grid.shape))\n",
    "    print(\"Batch Grids: {}\".format(batch_grids.shape))\n",
    "\n",
    "    return batch_grids\n",
    "\n",
    "\n",
    "def bilinear_sampler(input_img, x, y):\n",
    "    \"\"\"\n",
    "    如果想要测试是否正常，将仿射变换改为恒等变换，查看是输入和输出\n",
    "    \"\"\"\n",
    "    # grab dimensions\n",
    "    B, H, W, C = input_img.shape\n",
    "\n",
    "    # 原本x,y在[-1,1]之间，现在放缩到 [0, W/H]\n",
    "    x = ((x + 1.) * W) * 0.5\n",
    "    y = ((y + 1.) * H) * 0.5\n",
    "\n",
    "    # 获取到每个(x_i, y_i)相邻的4个坐标\n",
    "    x0 = np.floor(x).astype(np.int64)\n",
    "    x1 = x0 + 1\n",
    "    y0 = np.floor(y).astype(np.int64)\n",
    "    y1 = y0 + 1\n",
    "\n",
    "    # 确保不会超界  make sure it's inside img range [0, H] or [0, W]\n",
    "    x0 = np.clip(x0, 0, W-1)\n",
    "    x1 = np.clip(x1, 0, W-1)\n",
    "    y0 = np.clip(y0, 0, H-1)\n",
    "    y1 = np.clip(y1, 0, H-1)\n",
    "\n",
    "    # 取出对应的像素值\n",
    "    Ia = input_img[np.arange(B)[:,None,None], y0, x0]\n",
    "    Ib = input_img[np.arange(B)[:,None,None], y1, x0]\n",
    "    Ic = input_img[np.arange(B)[:,None,None], y0, x1]\n",
    "    Id = input_img[np.arange(B)[:,None,None], y1, x1]\n",
    "\n",
    "    # 依据双线性插值公式，计算deltas\n",
    "    wa = (x1-x) * (y1-y)\n",
    "    wb = (x1-x) * (y-y0)\n",
    "    wc = (x-x0) * (y1-y)\n",
    "    wd = (x-x0) * (y-y0)\n",
    "\n",
    "    # add dimension for addition\n",
    "    wa = np.expand_dims(wa, axis=3)\n",
    "    wb = np.expand_dims(wb, axis=3)\n",
    "    wc = np.expand_dims(wc, axis=3)\n",
    "    wd = np.expand_dims(wd, axis=3)\n",
    "\n",
    "    # compute output\n",
    "    out = wa*Ia + wb*Ib + wc*Ic + wd*Id\n",
    "\n",
    "    return out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 设置仿射矩阵\n",
    "\n",
    "- 设置仿射矩阵的值，对应不同的仿射变换\n",
    "- 调节仿射矩阵的batch维度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 2, 3)\n"
     ]
    }
   ],
   "source": [
    "B, H, W, C = input_img.shape\n",
    "\n",
    "# 设置仿射矩阵\n",
    "#M = np.array([[1., 0., 0.], [0., 1., 0.]]) 恒等变换\n",
    "M = np.array([[1., 0., 0.5], [0., 1., 0.]]) #平移\n",
    "#M = np.array([[0.707, -0.707, 0.], [0.707, 0.707, 0.]]) #旋转\n",
    "\n",
    "# 调整维度\n",
    "M = np.resize(M, (B, 2, 3))\n",
    "print (M.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 获取仿射矩阵采样点，采样\n",
    "\n",
    "- 获取仿射矩阵生成的采样点\n",
    "- 将x,y维度拆分\n",
    "- 使用双线性插值采样"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformation Matrices: (2, 2, 3)\n",
      "Sampling Grid: (2, 3, 160000)\n",
      "Batch Grids: (2, 400, 400, 2)\n"
     ]
    }
   ],
   "source": [
    "batch_grids = affine_grid_generator(H, W, M)\n",
    "\n",
    "x_s = batch_grids[:, :, :, 0:1].squeeze()\n",
    "y_s = batch_grids[:, :, :, 1:2].squeeze()\n",
    "\n",
    "out = bilinear_sampler(input_img, x_s, y_s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 查看变换图片"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = array_to_img(out[-1])\n",
    "out.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
